import pylab as plt
import sys
import os
import re
import numpy as np
from params import *

DIR = os.path.join(WORKDIR, 'gowalla', 'results')

data = {}
for file in os.listdir(DIR):
    if not file.startswith('sampled_distance'):
        continue
    root, ext = os.path.splitext(file)
    date = root.split('_')[-1]
    month = int(date.split('-')[1])
    print date,month
    
    MAX_HOP = 7
    for line in open(os.path.join(DIR,file)):
        print line
        tokens = map(int,line.split(','))
        hop,pos,neg = tokens[:3]
        #divide by 2 to avoid counting edges twice
        pos /= 2
        neg /= 2
        hop = min(MAX_HOP,hop)
        data.setdefault(month,{}).setdefault(hop,(0,0)) 
        data[month][hop] = map(sum, zip(data[month][hop],(pos,neg)))

pos_v, neg_v, imb_v = [],[],[]
lbl = []
for month in data:
    print month
    tot_pos,tot_neg = 0.0, 0.0
    for hop in data[month]:
        print hop, data[month][hop]
        pos,neg = data[month][hop]
        tot_pos += pos
        tot_neg += neg
    lbl.append(month)
    pos_v.append(tot_pos/2)
    neg_v.append(tot_neg/2)
    imb_v.append(tot_neg/tot_pos)

plt.figure()
plt.semilogy(pos_v,'kx-')
plt.semilogy(neg_v,'k+-')
plt.semilogy(imb_v,'ko-')
plt.legend(('Positive', 'Negative', 'Imbalance'))
plt.grid(True)
plt.savefig('all_imbalance.pdf')
plt.close()

#for month in sorted(data):
#    print ''
#    print month
#    pos_v, neg_v, imb_v = [],[],[]
#    ind = []
#    tot = 0
#    for hop in sorted(data[month]):
#        if hop < 2:
#            continue
#        pos,neg = data[month][hop]
#        print hop, pos,neg
#        pos_v.append(pos)
#        neg_v.append(neg)
#        imb_v.append(1.0*neg/(pos+1))
#        ind.append(hop)
#        tot += pos
#
#    print tot
#    width = 0.25       # the width of the bars
#    ind = np.array(ind)
#
#    plt.figure()
#    plt.clf()
#    plt.axes(FIG_AXES)
#
#    #rects1 = plt.bar(ind-width, pos_v, width, color='1.00',log=True,linewidth=2)
#    #rects2 = plt.bar(ind, neg_v, width, color='0.00',log=True,linewidth=2)
#    #rects3 = plt.bar(ind+width, imb_v, width, color='0.50',log=True,linewidth=2)
#    #plt.legend([rects1[0],rects2[0],rects3[0]], ('Positive','Negative','Imbalance'),
#    #        loc='lower left')
#
#
#    plt.semilogy(ind,pos_v,'kx-')
#    plt.semilogy(ind,neg_v,'k.-')
#    plt.semilogy(ind,imb_v,'k+-')
#    plt.legend(('Future linked pairs','Future unlinked pairs','Unlinked/linked ratio'),
#            loc='lower left')
#
#    ticks = ind
#    lbl = map(str,ind)
#    plt.xticks(ticks,lbl)
#    plt.xlabel('Social network distance')
#    [x1,x2,y1,y2] = plt.axis()
#    plt.axis([2,x2,1,y2])
#    plt.grid(True)
#
#    plt.savefig('hop_imbalance_%d.pdf'%month)
#    plt.close()
#

for month in sorted(data):
    print ''
    print month
    pos_v, neg_v, imb_v = [],[],[]
    ind = []
    tot = 0
    for hop in sorted(data[month]):
        if hop < 2:
            continue
        pos,neg = data[month][hop]
        print hop, pos,neg
        pos_v.append(pos)
        neg_v.append(neg)
        imb_v.append(1.0*pos/(pos+neg))
        ind.append(hop)
        tot += pos

    print tot
    width = 0.25       # the width of the bars
    ind = np.array(ind)


    #rects1 = plt.bar(ind-width, pos_v, width, color='1.00',log=True,linewidth=2)
    #rects2 = plt.bar(ind, neg_v, width, color='0.00',log=True,linewidth=2)
    #rects3 = plt.bar(ind+width, imb_v, width, color='0.50',log=True,linewidth=2)
    #plt.legend([rects1[0],rects2[0],rects3[0]], ('Positive','Negative','Imbalance'),
    #        loc='lower left')


    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.semilogy(ind,pos_v,'ko', ms=6)
    ticks = ind
    lbl = map(str,ind)
    #plt.xticks(ticks,lbl)
    plt.xlabel('Social network distance')
    plt.ylabel('New links')
    [x1,x2,y1,y2] = plt.axis()
    plt.axis([2,x2,y1,y2])
    plt.grid(True)
    plt.savefig('hop_new_links_%d.pdf'%month)
    plt.close()
    print pos_v


    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.semilogy(ind,imb_v,'ko', ms=6)
    ticks = ind
    lbl = map(str,ind)
    #plt.xticks(ticks,lbl)
    plt.xlabel('Social network distance')
    plt.ylabel('Link probability')
    [x1,x2,y1,y2] = plt.axis()
    plt.axis([2,x2,y1,y2])
    plt.grid(True)
    plt.savefig('hop_new_prob_%d.pdf'%month)
    plt.close()
    print imb_v

